"""
使用transformers库的翻译服务器
不依赖libsndfile.dll
"""

from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

app = Flask(__name__)
CORS(app)

print("正在加载翻译模型...")
model = M2M100ForConditionalGeneration.from_pretrained("facebook/m2m100_418M")
tokenizer = M2M100Tokenizer.from_pretrained("facebook/m2m100_418M")

def translate_text(text, src_lang="en", tgt_lang="zh"):
    tokenizer.src_lang = src_lang
    encoded_text = tokenizer(text, return_tensors="pt")
    generated_tokens = model.generate(**encoded_text, forced_bos_token_id=tokenizer.get_lang_id(tgt_lang))
    return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]

@app.route('/translate', methods=['POST'])
def translate():
    try:
        data = request.json
        text = data.get('text', '')
        
        if not text:
            return jsonify({'error': '文本为空'}), 400
        
        translation = translate_text(text)
        return jsonify({'translation': translation})
    
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health():
    return jsonify({
        'status': 'ready',
        'model': 'facebook/m2m100_418M'
    })

if __name__ == '__main__':
    print("=" * 60)
    print("Transformers翻译服务器启动中...")
    print("=" * 60)
    print("\n✓ 模型已加载")
    print("\n服务器地址: http://localhost:5000")
    print("翻译接口: POST http://localhost:5000/translate")
    print("健康检查: GET http://localhost:5000/health")
    print("\n按 Ctrl+C 停止服务器")
    print("=" * 60)
    
    app.run(host='0.0.0.0', port=5000, debug=False)